import numpy as np
import pandas as pd

np.random.seed(42)
n_samples = 500

# 特征生成
has_spam_keyword = np.random.choice(
    [0, 1], size=n_samples, p=[0.7, 0.3]
)  # 是否包含垃圾关键词
unknown_sender = np.random.choice(
    [0, 1], size=n_samples, p=[0.6, 0.4]
)  # 是否未知发件人
is_image_only = np.random.choice([0, 1], size=n_samples, p=[0.8, 0.2])  # 是否仅包含图片
length_norm = np.clip(
    np.random.normal(loc=0.5, scale=0.2, size=n_samples), 0, 1
)  # 长度归一化
punctuation_ratio = np.clip(np.random.beta(2, 5, size=n_samples), 0, 1)  # 标点符号比例

# 构建特征矩阵
X = np.stack(
    [has_spam_keyword, unknown_sender, is_image_only, length_norm, punctuation_ratio],
    axis=1,
)

# 构建标签（根据合理规则，不是完全随机）
y = []
for i in range(n_samples):
    if has_spam_keyword[i] and unknown_sender[i]:
        y.append(1)
    elif is_image_only[i] and length_norm[i] < 0.3:
        y.append(1)
    elif punctuation_ratio[i] > 0.6:
        y.append(1)
    else:
        y.append(
            np.random.choice([0, 1], p=[0.9, 0.1])
        )  # 有少量正常邮件也标记为垃圾（人为加入混淆）
y = np.array(y)

# 转为 DataFrame 显示
df = pd.DataFrame(
    X,
    columns=[
        "has_spam_keyword",
        "unknown_sender",
        "is_image_only",
        "length_norm",
        "punctuation_ratio",
    ],
)
df["label"] = y

print(df.head(10))

# 保存为 CSV 文件
df.to_csv("spam_dataset.csv", index=False, encoding="utf-8-sig")
print("数据已成功保存为 spam_dataset.csv")
